"""
日志工具模块，实现日志记录功能。
"""
import os
import logging
import time
import sys
from typing import Optional, Union, Dict, Any, List
from enum import Enum

class LogLevel(Enum):
    """日志级别枚举"""
    DEBUG = logging.DEBUG
    INFO = logging.INFO
    WARNING = logging.WARNING
    ERROR = logging.ERROR
    CRITICAL = logging.CRITICAL

def setup_logger(name: str, 
                level: LogLevel = LogLevel.INFO,
                log_file: Optional[str] = None,
                console_output: bool = True,
                file_format: str = '%(asctime)s - %(name)s - %(levelname)s - %(message)s') -> logging.Logger:
    """
    设置日志记录器
    
    Args:
        name: 日志记录器名称
        level: 日志级别
        log_file: 日志文件路径
        console_output: 是否输出到控制台
        file_format: 文件格式
    
    Returns:
        日志记录器实例
    """
    logger = logging.getLogger(name)
    logger.setLevel(level.value)
    
    # 清除已有的处理器
    for handler in logger.handlers[:]:
        logger.removeHandler(handler)
    
    # 创建格式器
    formatter = logging.Formatter(file_format)
    
    # 添加文件处理器
    if log_file:
        # 确保目录存在
        os.makedirs(os.path.dirname(os.path.abspath(log_file)), exist_ok=True)
        
        file_handler = logging.FileHandler(log_file)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
    
    # 添加控制台处理器
    if console_output:
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)
    
    return logger

class MonitorLogger:
    """监控日志记录器"""
    
    def __init__(self, 
                name: str = "colo_monitor",
                level: LogLevel = LogLevel.INFO,
                log_dir: Optional[str] = None):
        """
        初始化监控日志记录器
        
        Args:
            name: 日志记录器名称
            level: 日志级别
            log_dir: 日志目录
        """
        self.name = name
        self.level = level
        self.log_dir = log_dir
        
        if log_dir:
            os.makedirs(log_dir, exist_ok=True)
            timestamp = time.strftime("%Y%m%d-%H%M%S")
            log_file = os.path.join(log_dir, f"{name}_{timestamp}.log")
        else:
            log_file = None
        
        self.logger = setup_logger(name, level, log_file)
    
    def debug(self, message: str, **kwargs) -> None:
        """
        记录调试级别的日志
        
        Args:
            message: 日志消息
            **kwargs: 其他参数
        """
        self.logger.debug(message, **kwargs)
    
    def info(self, message: str, **kwargs) -> None:
        """
        记录信息级别的日志
        
        Args:
            message: 日志消息
            **kwargs: 其他参数
        """
        self.logger.info(message, **kwargs)
    
    def warning(self, message: str, **kwargs) -> None:
        """
        记录警告级别的日志
        
        Args:
            message: 日志消息
            **kwargs: 其他参数
        """
        self.logger.warning(message, **kwargs)
    
    def error(self, message: str, **kwargs) -> None:
        """
        记录错误级别的日志
        
        Args:
            message: 日志消息
            **kwargs: 其他参数
        """
        self.logger.error(message, **kwargs)
    
    def critical(self, message: str, **kwargs) -> None:
        """
        记录严重错误级别的日志
        
        Args:
            message: 日志消息
            **kwargs: 其他参数
        """
        self.logger.critical(message, **kwargs)
    
    def log_param(self, name: str, value: Any, step: Optional[int] = None) -> None:
        """
        记录参数信息
        
        Args:
            name: 参数名称
            value: 参数值
            step: 当前步数
        """
        step_info = f" (step {step})" if step is not None else ""
        self.info(f"Parameter {name}: {value}{step_info}")
    
    def log_metric(self, name: str, value: Any, step: Optional[int] = None) -> None:
        """
        记录指标信息
        
        Args:
            name: 指标名称
            value: 指标值
            step: 当前步数
        """
        step_info = f" (step {step})" if step is not None else ""
        self.info(f"Metric {name}: {value}{step_info}")
    
    def log_anomaly(self, anomaly_type: str, parameter_name: str, value: Any, threshold: Any, step: Optional[int] = None) -> None:
        """
        记录异常信息
        
        Args:
            anomaly_type: 异常类型
            parameter_name: 参数名称
            value: 异常值
            threshold: 阈值
            step: 当前步数
        """
        step_info = f" (step {step})" if step is not None else ""
        self.warning(f"Anomaly detected: {anomaly_type} in {parameter_name}, value={value}, threshold={threshold}{step_info}")
    
    def log_model_info(self, model: Any) -> None:
        """
        记录模型信息
        
        Args:
            model: 模型实例
        """
        self.info(f"Model type: {type(model).__name__}")
        
        # 统计参数数量
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        
        self.info(f"Total parameters: {total_params:,}")
        self.info(f"Trainable parameters: {trainable_params:,}")
        self.info(f"Non-trainable parameters: {total_params - trainable_params:,}")
    
    def log_memory_stats(self, stats: Dict[str, Any]) -> None:
        """
        记录内存统计信息
        
        Args:
            stats: 内存统计字典
        """
        # 系统内存
        if "system" in stats:
            sys_mem = stats["system"]
            total_gb = sys_mem["total"] / (1024 ** 3)
            used_gb = sys_mem["used"] / (1024 ** 3)
            self.info(f"System memory: {used_gb:.2f}GB / {total_gb:.2f}GB ({sys_mem['percent']}%)")
        
        # 进程内存
        if "process" in stats:
            proc_mem = stats["process"]
            rss_gb = proc_mem["rss"] / (1024 ** 3)
            vms_gb = proc_mem["vms"] / (1024 ** 3)
            self.info(f"Process memory: RSS={rss_gb:.2f}GB, VMS={vms_gb:.2f}GB")
        
        # CUDA内存
        if "cuda" in stats:
            for gpu_id, gpu_stats in stats["cuda"].items():
                allocated_mb = gpu_stats["allocated"] / (1024 ** 2)
                reserved_mb = gpu_stats["reserved"] / (1024 ** 2)
                self.info(f"CUDA {gpu_id}: Allocated={allocated_mb:.2f}MB, Reserved={reserved_mb:.2f}MB") 